#! /usr/bin/python2.5
# This file is part of SyncStar, free software released under GPL v2.
# Google AppEngine Python 2.5 source.

from google.appengine.api import apiproxy_stub_map

DATASTORE_RPCS_WITH_COST = ('Commit', 'Delete', 'Put', 'Touch')

DATASTORE_RPCS = (
    'AddActions', 'AllocateIds', 'BeginTransaction', 'Commit', 'Count',
    'Delete', 'Get', 'GetIndices', 'Next', 'Put', 'Rollback',
    'RunCompiledQuery', 'RunQuery', 'Touch')


class DbCost(object):
  __slots__ = ('index_writes', 'index_write_bytes', 'entity_writes',
               'entity_write_bytes', 'cpu_usage_mcycles',
               'requested_entity_puts', 'requested_entity_deletes') + (
                  tuple(DATASTORE_RPCS))

  def __init__(self, other=None):
    if other is None:
      for name in self.__slots__:
        setattr(self, name, 0)
    else:
      for name in self.__slots__:
        setattr(self, name, getattr(other, name))

  def __isub__(self, other):
    for name in self.__slots__:
      setattr(self, name, getattr(self, name) - getattr(other, name))
    return self

  def __sub__(self, other):
    result = self.copy()
    result -= other
    return result

  def copy(self):
    return type(self)(self)

  def add_cost(self, cost):
    # Cost fields are from datastore_pb.py .
    assert cost
    self.index_writes += cost.index_writes()
    self.index_write_bytes += cost.index_write_bytes()
    self.entity_writes += cost.entity_writes()
    self.entity_write_bytes += cost.entity_write_bytes()
    if cost.has_commitcost():
      self.requested_entity_puts += cost.commitcost().requested_entity_puts()
      self.requested_entity_deletes += (
          cost.commitcost().requested_entity_deletes())

  def add_rpc(self, rpc):
    assert rpc
    self.cpu_usage_mcycles += rpc.cpu_usage_mcycles  # Megacycles on an 1.2GHz.

  def add_call(self, call):
    assert isinstance(call, str)
    if call in DATASTORE_RPCS:
      setattr(self, call, getattr(self, call) + 1)

  def format_short(self):
    output = []
    for name in sorted(self.__slots__):
      value = getattr(self, name)
      if value:  # Discard zero values.
        output.append('%s:%s' % (name, value))
    if output:
      return ' '.join(output)
    else:
      return 'no_cost'


global_db_cost = None


def _db_count_hook(service, call, request, response, rpc, *args, **kwargs):
  # Typical call names for service == 'datastore_v3':
  #     'Get', 'RunQuery', 'BeginTransaction', 'Put', 'Commit'.
  # All call names for service == 'datastore_v3' (from datastore_pb.py,
  # datastore_rpc.py, datastore_query.py and remote_api_services.py):
  #     'AddActions', 'AllocateIds', 'BeginTransaction', 'Commit', 'Count',
  #     'Delete', 'Get', 'GetIndices', 'Next', 'Put', 'Rollback',
  #     'RunCompiledQuery', 'RunQuery', 'Touch'.
  # All call names for service == 'datastore_v3' which have a cost
  # (from https://gist.github.com/715284):
  #     'Commit', 'Delete', 'Put', 'Touch'.
  # Typical call names for service == 'user':
  #     'CreateLoginURL', 'CreateLogoutURL'.
  # Unfortunately this doesn't count the number of datastore operations.
  if service == 'datastore_v3':
    if call in DATASTORE_RPCS:
      if call in DATASTORE_RPCS_WITH_COST and response.has_cost():
        global_db_cost.add_cost(response.cost())
      global_db_cost.add_rpc(rpc)
      global_db_cost.add_call(call)


def register_db_count_hook():
  global global_db_cost
  if global_db_cost is None:
    global_db_cost = DbCost()
    apiproxy_stub_map.apiproxy.GetPostCallHooks().Append(
       'db_count_hook', _db_count_hook, 'datastore_v3')

